#ifndef _SKIPLIST_H_
#define _SKIPLIST_H_

#include <vector>
#include <memory>

// 非线程安全的跳表实现
template <typename Key, typename Value>
class SkipList {
public:
    SkipList();    

    SkipList(const SkipList&) = delete;
    SkipList& operator=(const SkipList&) = delete;
    
    void Insert(const Key& k, const Value& v);

    bool Find(const Key& k, Value* v);

    void Delete(const Key& k);

private:
    static constexpr int kMaxHeight = 24;

    struct SkipListNode; 

    typedef std::shared_ptr<SkipListNode> SkipNodePtr;
    typedef std::vector<SkipNodePtr>      SkipNodePtrVec;

    SkipNodePtr NewNode(const Key& k, const Value& v, int height);

    int RandomHeight();

    inline int GetMaxHeight() { return max_height; }

private:
    int max_height;   // 当前跳表中的最高高度
    SkipNodePtr head; // 虚拟头结点
};


template <typename Key, typename Value>
struct SkipList<Key, Value>::SkipListNode {
    SkipListNode(const Key& k, const Value& v, int height) 
    : key(k), value(v), next(height, nullptr)
    {}
    Key const key;
    Value value;
    SkipNodePtrVec next;        
};

template <typename Key, typename Value>
typename SkipList<Key, Value>::SkipNodePtr 
SkipList<Key, Value>::NewNode(const Key& k, const Value& v, int height) {
    return std::make_shared<SkipListNode>(k, v, height);
}

template <typename Key, typename Value>
int SkipList<Key, Value>::RandomHeight() {
    int height = 1;
    while (height < kMaxHeight && rand() % 2) {
        ++height;
    }
    return height;
}

template <typename Key, typename Value>
SkipList<Key, Value>::SkipList() 
    : max_height(1),
      head(NewNode(Key(), Value(), kMaxHeight)) {
    srand(time(nullptr));
}

template <typename Key, typename Value>
void SkipList<Key, Value>::Insert(const Key& k, const Value& v) {
    SkipNodePtr prev[kMaxHeight];
    SkipNodePtr curr = head;
    
    int level = GetMaxHeight() - 1;
    while (true) {
        SkipNodePtr next = curr->next[level];
        if (next != nullptr && next->key < k) {
            curr = next;
        } else {
            prev[level] = curr;
            if (level == 0) {
                curr = next;
                break;
            } else {
                level--;
            }
        }
    }
    
    // 找到相同 key 的 Node 就更新其 value
    if (curr != nullptr && curr->key == k) {
        curr->value = v;
        return;
    }

    int height = RandomHeight();
    if (height > max_height) {
        for (int i = max_height; i < height; i++) {
            prev[i] = head;
        }
        max_height = height;
    }

    SkipNodePtr node = NewNode(k, v, height);
    for (int i = 0; i < height; i++) {
        node->next[i] = prev[i]->next[i];
        prev[i]->next[i] = node;
    }
}

template <typename Key, typename Value>
bool SkipList<Key, Value>::Find(const Key& k, Value* v) {
    SkipNodePtr curr = head;
    int level = GetMaxHeight() - 1;
       
    while (true) {
        SkipNodePtr next = curr->next[level];
        if (next != nullptr) {
            if (next->key < k) {
                curr = next;
            } else if (next->key == k) {
                *v = next->value;
                return true;
            } else {
                if (level == 0) return false;
                level--;
            }
        } else {
            if (level == 0) return false;
            level--;
        }
    }
    return false;
}

template <typename Key, typename Value>
void SkipList<Key, Value>::Delete(const Key& k) {
    SkipNodePtr prev[kMaxHeight];
    SkipNodePtr curr = head;

    int level = GetMaxHeight() - 1;
    while (true) {
        SkipNodePtr next = curr->next[level];
        if (next != nullptr && next->key < k) {
            curr = next;
        } else {
            prev[level] = curr;
            if (level == 0) {
                curr = next;
                break;
            }
            level--;
        }
    }

    if (curr == nullptr || curr->key != k) return;
    
    for (int i = GetMaxHeight() - 1; i >= 0; i--) {
        if (prev[i]->next[i] == curr)
            prev[i]->next[i] = curr->next[i];
    }   
}

#endif